from cProfile import label
import os
import pprint
import shutil

import json
from tqdm import tqdm, trange
import numpy as np
import torch
import torch.nn.functional as F
import pandas as pd              #biran 

from common.utils.load import smart_load_model_state_dict
from common.trainer import to_cuda
from common.utils.create_logger import create_logger
from fakenews.data.build import make_dataloader
from fakenews.modules import *


@torch.no_grad()
def test_net(args, config, ckpt_path=None, save_path=None, save_name=None):
    print('test net...')
    pprint.pprint(args)
    pprint.pprint(config)
    device_ids = [int(d) for d in config.GPUS.split(',')]
    # os.environ['CUDA_VISIBLE_DEVICES'] = config.GPUS

    torch.backends.cudnn.enabled = False   # 禁止CuDNN使用非确定性算法
    torch.backends.cudnn.deterministic = True    # 避免网络前馈结果波动
    torch.backends.cudnn.benchmark = False    # 禁止寻找最优卷积算法以花费时间

    if ckpt_path is None:
        _, train_output_path = create_logger(config.OUTPUT_PATH, args.cfg, config.DATASET.TRAIN_IMAGE_SET,
                                             split='train')
        model_prefix = os.path.join(train_output_path, config.MODEL_PREFIX)
        ckpt_path = '{}-best.model'.format(model_prefix)
        print('Use best checkpoint {}...'.format(ckpt_path))
    if save_path is None:
        logger, test_output_path = create_logger(config.OUTPUT_PATH, args.cfg, config.DATASET.TEST_IMAGE_SET,
                                                 split='test')
        save_path = test_output_path
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    shutil.copy2(ckpt_path, os.path.join(save_path, '{}_test_ckpt_{}.model'.format(config.MODEL_PREFIX, config.DATASET.TASK)))

    # get network
    model = eval(config.MODULE)(config)
    if len(device_ids) > 1:
        model = torch.nn.DataParallel(model, device_ids=device_ids).cuda()
    else:
        torch.cuda.set_device(device_ids[0])
        model = model.cuda()
    checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
    smart_load_model_state_dict(model, checkpoint['state_dict'])

    # loader
    test_loader = make_dataloader(config, mode='test', distributed=False)
    test_dataset = test_loader.dataset
    test_database = test_dataset.database
    
    #br  加载csv文件中的label
    label = []
    filename = 'data/PhemeData/test_phemedata_filter15.csv'       #tweets_fn    #config.DATASET.TEST_IMAGE_SET
    data = pd.read_csv(filename, sep=',') 
    data=data['label'].values.tolist()       #取出label这一列
    label = [int(u in 'real') for u in data]    # label ={'real':1,'fake':0}
    
    print("len(data):",len(data))
    drop_count = len(data) % model.batch_size
    if drop_count != 0: 
        data = data[0:len(data) - drop_count]
        print("{} were dropped to align with {}".format(drop_count, model.batch_size))
    
    # test
    q_ids = []
    answer_ids = []
    model.eval()
    cur_id = 0
    for nbatch, batch in zip(trange(len(test_loader)), test_loader):
    # for nbatch, batch in tqdm(enumerate(test_loader)):
        bs = test_loader.batch_sampler.batch_size if test_loader.batch_sampler is not None else test_loader.batch_size
        #q_ids.extend([test_database[id]['question_id'] for id in range(cur_id, min(cur_id + bs, len(test_database)))])
        batch = to_cuda(batch)
        output = model(*batch)
        answer_ids.extend(output['label_logits'].argmax(dim=1).detach().cpu().tolist())
        cur_id += bs
            
    result = [{'question_id': q_id, 'answer': test_dataset.answer_vocab[a_id]} for q_id, a_id in zip(q_ids, answer_ids)]
    
    # br 计算acc
    count = 0
    for i in range(len(test_database) - 1):
        if answer_ids[i] == label[i]:
            count += 1
        else:
            continue
    result = count / len(test_database)
    print('Accuracy_score =', result)

    cfg_name = os.path.splitext(os.path.basename(args.cfg))[0]
    result_json_path = os.path.join(save_path, '{}_vqa2_{}.json'.format(cfg_name if save_name is None else save_name,
                                                                        config.DATASET.TEST_IMAGE_SET))
    with open(result_json_path, 'w') as f:
        json.dump(answer_ids, f)
    print('result json saved to {}.'.format(result_json_path))
    return result_json_path
